package com.weijian.config;

import com.knuddels.jtokkit.api.EncodingType;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepository;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.image.ImageModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiImageModel;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * 模型初始化全局配置（这里向量数据库只有唯一es，所以没有初始配置，默认就是）
 */
@Configuration
@RequiredArgsConstructor
public class ChatClientConfig {

    final ZhiPuAiChatModel zhiPuAiChatModel;

    final OpenAiChatModel openAiChatModel;

    final OllamaEmbeddingModel ollamaEmbeddingModel;

    final OpenAiEmbeddingModel openAiEmbeddingModel;

    final ZhiPuAiImageModel zhiPuAiImageModel;

//    final JdbcChatMemoryRepository chatMemoryRepository;


    /**
     * 对话模型的全局初始化配置
     */
    @Bean
    public ChatClient chatClient() {
        return ChatClient.builder(zhiPuAiChatModel)
                .build();
    }

//    @Bean
//    public ChatMemory chatMemory() {
//        return MessageWindowChatMemory.builder()
//                .chatMemoryRepository(chatMemoryRepository)
//                .maxMessages(10)
//                .build();
//    }

    /**
     * 文本嵌入模型的全局初始化配置
     */
    @Bean
    public EmbeddingModel embeddingModel() {
        return ollamaEmbeddingModel;
    }

    @Bean
    public ImageModel imageModel() {
        return zhiPuAiImageModel;
    }


    /**
     * 文本嵌入的批量策略，单次最大长度设置为20000，最大占比设置为0.1（默认情况8191，文本长度有可能超出）
     */
    @Bean
    public BatchingStrategy customTokenCountBatchingStrategy() {
        return new TokenCountBatchingStrategy(
                EncodingType.CL100K_BASE,
                30000,
                0.1
        );
    }
}
